import * as tf from '@tensorflow/tfjs';
import { LSTM } from '../models/lstmModel';
import { Transformer } from '../models/transformerModel';
import { preprocessData } from './dataPreprocessor';

export async function tuneHyperparameters(data, modelType, maxSequenceLength, logCallback, progressCallback) {
  // Define search space
  const paramsSpace = {
    sequenceLength: [5, 10, 15, 20],
    units: [32, 64, 128],
    learningRate: [0.001, 0.0001],
    batchSize: [16, 32, 64],
    epochs: [30, 50, 100],
    ...(modelType === 'transformer' && {
      dModel: [32, 64, 128],
      numHeads: [2, 4, 8]
    })
  };
  
  // Generate random combinations
  const trials = 10;
  const trialsParams = [];
  
  for (let i = 0; i < trials; i++) {
    const params = {
      sequenceLength: paramsSpace.sequenceLength[
        Math.floor(Math.random() * paramsSpace.sequenceLength.length)
      ],
      units: paramsSpace.units[
        Math.floor(Math.random() * paramsSpace.units.length)
      ],
      learningRate: paramsSpace.learningRate[
        Math.floor(Math.random() * paramsSpace.learningRate.length)
      ],
      batchSize: paramsSpace.batchSize[
        Math.floor(Math.random() * paramsSpace.batchSize.length)
      ],
      epochs: paramsSpace.epochs[
        Math.floor(Math.random() * paramsSpace.epochs.length)
      ]
    };
    
    if (modelType === 'transformer') {
      params.dModel = paramsSpace.dModel[
        Math.floor(Math.random() * paramsSpace.dModel.length)
      ];
      params.numHeads = paramsSpace.numHeads[
        Math.floor(Math.random() * paramsSpace.numHeads.length)
      ];
    }
    
    trialsParams.push(params);
  }
  
  // Evaluate each combination
  let bestParams = null;
  let bestValAcc = 0;
  
  for (let i = 0; i < trialsParams.length; i++) {
    const params = trialsParams[i];
    progressCallback((i / trialsParams.length) * 100);
    
    logCallback(`Trial ${i + 1}/${trialsParams.length}: Testing params ${JSON.stringify(params)}`);
    
    try {
      const { X, y, numClasses } = preprocessData(data, params.sequenceLength);
      
      let model;
      if (modelType === 'lstm') {
        model = LSTM(numClasses, params.sequenceLength, params.units, params.learningRate);
      } else {
        model = Transformer(
          numClasses, 
          params.sequenceLength, 
          params.dModel, 
          params.numHeads, 
          params.learningRate
        );
      }
      
      const history = await model.fit(X, y, {
        epochs: params.epochs,
        batchSize: params.batchSize,
        validationSplit: 0.2,
        verbose: 0
      });
      
      const valAcc = history.history.val_acc.slice(-1)[0];
      logCallback(`Trial ${i + 1} validation accuracy: ${valAcc.toFixed(4)}`);
      
      if (valAcc > bestValAcc) {
        bestValAcc = valAcc;
        bestParams = params;
        logCallback(`New best params found with val_acc: ${valAcc.toFixed(4)}`);
      }
      
      // Clean up memory
      tf.dispose([X, y, model]);
    } catch (error) {
      logCallback(`Error in trial ${i + 1}: ${error.message}`);
    }
  }
  
  progressCallback(100);
  logCallback(`Hyperparameter tuning completed. Best val_acc: ${bestValAcc.toFixed(4)}`);
  
  return bestParams;
}